import cv2
import numpy as np
import onnxruntime

import sys
import os

conf_thres = 0.5
nms_thres = 0.3
num_anchor = 3  # number of anchor
class_head = 1  # number of class
anchor = np.array([142, 110, 192, 243, 459, 401, 36, 75, 76, 55, 72, 146, 12, 16, 19, 36, 40, 28])
    
class Bbox:
    def __init__(self, box, score, label, status=False):
        self.box = box
        self.score = score
        self.label = label
        self.status = status

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def plotresult(img: np.ndarray, result: list):
    for bbox in result:
        cv2.rectangle(img, (bbox.box[3], bbox.box[2]), (bbox.box[1], bbox.box[0]), (0, 0, 255), thickness=1)
    cv2.imwrite("result.jpg", img)

def iou(bbox1:Bbox, bbox2:Bbox):
    x1_min, y1_min, x1_max, y1_max = bbox1.box

    x2_min, y2_min, x2_max, y2_max = bbox2.box
    
    inter_x_min = max(x1_min, x2_min)
    inter_y_min = max(y1_min, y2_min)
    inter_x_max = min(x1_max, x2_max)
    inter_y_max = min(y1_max, y2_max)
    
    inter_area = max(0, inter_x_max - inter_x_min) * max(0, inter_y_max - inter_y_min)
    
    area1 = (x1_max - x1_min) * (y1_max - y1_min)
    area2 = (x2_max - x2_min) * (y2_max - y2_min)
    
    union_area = area1 + area2 - inter_area
    
    iou = inter_area / union_area if union_area != 0 else 0
    
    return iou

# preprocess -> resize + pad ---> keep the radio of length to width
def preprocess(src_img, size=[640, 640]):
    image = cv2.imread(src_img)
    iw, ih = image.shape[1], image.shape[0]
    w, h = size

    scale = min(w / iw, h / ih)
    nw = int(iw * scale)
    nh = int(ih * scale)

    resized_image = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_CUBIC)

    new_image = np.full((h, w, 3), (128, 128, 128), dtype=np.uint8)

    x_offset = (w - nw) // 2
    y_offset = (h - nh) // 2

    new_image[y_offset:y_offset+nh, x_offset:x_offset+nw] = resized_image
    new_image = np.expand_dims(np.transpose(new_image, (2, 0, 1)), 0) / 255

    return new_image.astype(np.float32), size, [ih, iw]

# infer by onnxruntime
def infer(img, onnx_path):
    providers = ['CPUExecutionProvider']
    onnx_model = onnxruntime.InferenceSession(onnx_path, providers=providers)
    onnx_input = {onnx_model.get_inputs()[0].name: img}
    outputs = onnx_model.run(None, onnx_input)
    return outputs


# postprocess --> decode
"""
output:[tensor:1X255X20X20,tensor:1X255X40X40,tensor:1X255X80X80]
"""
def postprocess(outputs, input_shape, img_shape):
    img_shape = [img_shape[1], img_shape[0]]
    img_shape = np.array(img_shape)
    input_shape = np.array(input_shape)
    new_shape = np.round(img_shape * np.min(input_shape/img_shape))
    offset  = (input_shape - new_shape)/2./input_shape
    scale   = input_shape/new_shape

    for i, output in enumerate(outputs):  # should compute in model
        outputs[i] = sigmoid(output)
    # step1 find object
    result_obj = []
    for index, output in enumerate(outputs):
        width_head = output.shape[2]
        height_head = output.shape[3]
        anchor_scaler_x = input_shape[0] / width_head
        anchor_scaler_y = input_shape[1] / height_head
        output = output.reshape((-1, 1))

        for w in range(width_head):
            for h in range(height_head):
                for a in range(num_anchor):
                    anchor_x = anchor[index * 6 + 2 * a] / anchor_scaler_x
                    anchor_y = anchor[index * 6 + 2 * a + 1] / anchor_scaler_y
                    exit_p = output[a * (class_head + 5) * width_head * height_head + 4 * width_head * height_head + w * height_head + h, 0]
                    max_value = conf_thres
                    max_index = -1

                    for k in range(5, class_head + 5):  # should obtain max value(compute in model)
                        score = exit_p * output[a * (class_head + 5) * width_head * height_head + k * width_head * height_head + w * height_head + h, 0]
                        if score > max_value:
                            max_value = score
                            max_index = k
                    if max_index != -1:
                        x_obj = output[a * (class_head + 5) * width_head * height_head + 0 * width_head * height_head + w * height_head + h, 0]
                        y_obj = output[a * (class_head + 5) * width_head * height_head + 1 * width_head * height_head + w * height_head + h, 0]
                        w_obj = output[a * (class_head + 5) * width_head * height_head + 2 * width_head * height_head + w * height_head + h, 0]
                        h_obj = output[a * (class_head + 5) * width_head * height_head + 3 * width_head * height_head + w * height_head + h, 0]
                        
                        x_obj = (x_obj * 2 - 0.5 + h) / width_head
                        y_obj = (y_obj * 2 - 0.5 + w) / height_head
                        w_obj = ((w_obj * 2) ** 2 * anchor_x) / width_head
                        h_obj = ((h_obj * 2) ** 2 * anchor_y) / height_head
                        
                        box_y = (y_obj - offset[1]) * scale[1]
                        box_x = (x_obj - offset[0]) * scale[0]
                        box_h = (h_obj) * scale[1]
                        box_w = (w_obj) * scale[0]
                        
                        top = (box_y - box_h / 2) * img_shape[1]
                        left = (box_x - box_w / 2) * img_shape[0]
                        bottom = (box_y + box_h / 2) * img_shape[1]
                        right = (box_x + box_w / 2) * img_shape[0]
                        
                        top = max(0, np.floor(top).astype('int32'))
                        left = max(0, np.floor(left).astype('int32'))
                        bottom = min(img_shape[1], np.floor(bottom).astype('int32'))
                        right = min(img_shape[0], np.floor(right).astype('int32'))
                        result_obj.append(Bbox([top, left, bottom, right], max_value, int(max_index - 5)))
    result_obj_nms = []
    # step2 nms
    result_obj_sorted = sorted(result_obj, key=lambda x: x.score, reverse=True)

    for i in range(len(result_obj_sorted)):
        if result_obj_sorted[i].status:
            continue
        result_obj_nms.append(result_obj_sorted[i])
        result_obj_sorted[i].status = True
        for j in range(i + 1, len(result_obj_sorted)):
            if result_obj_sorted[i].label != result_obj_sorted[j].label:
                continue
            if iou(result_obj_sorted[i], result_obj_sorted[j]) > nms_thres:
                result_obj_sorted[j].status = True
    return result_obj_nms

if __name__ == '__main__':
    img_file = "../deploy_tensorrt/img/street.jpg"
    onnx_path = "../deploy_tensorrt/model/models.onnx"
    
    img, input_shape, img_shape = preprocess(img_file)
    
    outputs = infer(img, onnx_path)
    
    result = postprocess(outputs, input_shape, img_shape)

    img = cv2.imread(img_file)
    plotresult(img, result)
